#!/usr/bin/python -S
"""
tnet_multi_test.py: Tests for tnet_multi.py
"""

__author__ = 'Andy Chu'


import os
import sys
import unittest

import jsontemplate

_test_cases = []
_counter = 0

def Test(**kwargs):
  global _counter
  # Test case can be named, or it gets a number as its name
  if 'name' not in kwargs:
    kwargs['name'] = str(_counter)
  _test_cases.append(kwargs)
  _counter += 1


Test(
  py = "None",
  js = "null",
  R = "NA",
  r_check = 'is.na(d)')

Test(
  name = "integer",
  py = "4",
  js = "4",
  R = "4L")

Test(
  name = "float",
  py = "5.9",
  js = "5.9",
  R = "5.9")

Test(
  name = "bool",
  py = "True",
  js = "true",
  R = "T")

Test(
  name = "empty-bytes",
  py = "u''",
  R = "''")

Test(
  name = "string",
  py = "u'hello'",
  js = "'hello'",
  R = "'hello'")

Test(
  name = "empty-string",
  py = "u''",
  js = "''",
  R = "''")

# TODO: Should R turn $ and , into R strings?  It has a "raw" type, but that is
# not useful for a lot of things.
Test(
  name = "unicode",
  py = "u'mu = \\u00B5'",
  js = "'hello'",
  R = "'mu = \\u00B5'")

# TODO: This exposes the bug in tnet.R - tnet.load.prefix.  Argh.
Test(
  name = "array-of-unicode",
  py = "[u'mu = \\u00B5', u'\\u00B6']",
  js = "'hello'",
  R = "c('mu = \\u00B5', '\\u00B6')")

Test(
  name = "empty-array",
  py = "[]",
  js = "[]",
  R = "list()",
  r_check = 'length(d) == 0',
  )

Test(
  name = "array",
  py = "[1]",
  js = "[1]",
  R = "list(1L)",  # use a list without names
  r_check = 'length(d) == 1 && d[[1]] == 1L')

# These array-of-* test the R vectorization optimization in tnet.dump.
Test(
  name = "array-of-int",
  py = "[1, 2]",
  js = "[1, 2]",
  R = "c(1L, 2L)",  # vector
  r_check = 'all(d == list(1L, 2L))')

Test(
  name = "array-of-float",
  py = "[1.2, 2.3]",
  js = "[1.2, 2.3]",
  R = "c(1.2, 2.3)",
  r_check = 'all(d == c(1.2, 2.3))')

Test(
  name = "array-of-bool",
  py = "[True, False]",
  js = "[true, false]",
  R = "c(T, F)",
  r_check = 'all(d == c(T, F))')

Test(
  name = "array-of-char",
  py = "[u'ab', u'cd']",
  js = "['ab', 'cd']",
  R = "c('ab', 'cd')",
  r_check = "all(d == c('ab', 'cd'))")

Test(
  name = "array-of-null",
  py = "[None]",
  js = "[null]",
  R = "list(NA)",
  r_check = 'is.list(d) && length(d) == 1')

Test(
  name = "array2",
  py = "[1,2]",
  js = "[1,2]",
  R = "list(1L, 2L)",
  r_check = 'all(d == c(1L, 2L))')

Test(
  name = "mixed-array",
  py = "[1, 2, False]",
  js = "[1, 2, false]",
  R = "list(1L, 2L, FALSE)",
  r_check = 'is.list(d) && d[[1]] == 1L && d[[2]] == 2L && d[[3]] == FALSE')

Test(
  name = "mixed-array2",
  py = "[1, 2, u'hi']",
  js = "[1, 2, 'hi']",
  R = "list(1L, 2L, 'hi')",
  r_check = "is.list(d) && d[[1]] == 1L && d[[2]] == 2L && d[[3]] == 'hi'")

# BUG: list() serializes as [] -- should we have a special empty.object() value?
# can't distinguish between {} and [] otherwise
Test(
  name = "empty-dict",
  py = "{}",
  js = "{}",
  R = "list()",
  r_check = 'length(d) == 0',
  )

Test(
  name = "dict1",
  py = "{u'a': 1}",
  js = "{a: 1}",
  R = "list(a=1L)",
  r_check = 'length(d) == 1 && d$a == 1L')

Test(
  name = "dict2",
  py = "{u'a': 1, u'b': 2}",
  js = "{a: 1, b: 2}",
  R = "list(a=1L, b=2L)",
  r_check = 'length(d) == 2 && d$a == 1L && d$b == 2L')

Test(
  name = "dict-of-array",
  py = "{u'a': [1,2,3]}",
  js = "{a: [1,2,3]}",
  R = "list(a=list(1L,2L,3L))",
  r_check = 'length(d) == 1 && all(d$a == c(1L, 2L, 3L))')

Test(
  name = "array-of-dict",
  py = "[{u'a': 1}, {u'a': 2, u'b': 3}]",
  js = "[{a: 1}, {a: 2, b: 3}]",
  R = "list(list(a=1L), list(a=2L, b=3L))",
  r_check = 'length(d) == 2 && length(d[[1]]) == 1 && length(d[[2]]) == 2')

Test(
  name = "array-of-array",
  py = "[[1,2], [3,4,5]]",
  js = "[[1,2], [3,4,5]]",
  R = "list(list(1L, 2L), list(3L, 4L, 5L))",
  r_check = 'length(d) == 2 && length(d[[1]]) == 2 && length(d[[2]]) == 3')

Test(
  name = "array-of-singleton-array",
  py = "[[7], [8]]",
  js = "[[7], [8]]",
  R = "list(list(7L), list(8L))",
  r_check = 'length(d) == 2 && length(d[[1]]) == 1 && length(d[[2]]) == 1')

# iteration order is different in R
Test(
  name = "dict-of-dict",
  py = "{u'a': {u'b': 2, u'c': 3}}",
  js = '{a: {b: 2, c: 3}}',
  R = "list(a=list(b=2L, c=3L))",
  r_check = 'length(d) == 1 && length(d$a) == 2 && '
            'd$a$b == 2L && d$a$c == 3L')


PY_TEST = jsontemplate.Template("""\
#!/usr/bin/python

import sys
import time

sys.path.append("{tnet_dir}/python")
import tnet

def DumpToFile(filename, data):
  f = open(filename, 'w')
  f.write(tnet.dumps(data))
  # reduce noise from diff -r
  f.write('\\n')
  f.close()


def DumpAll():
{.repeated section cases}
  DumpToFile("py/{name}.tnet", {py})
{.end}


def LoadFromFile(filename):
  f = open(filename)
  d = tnet.load(f)
  f.close()
  return d


def LoadAll():
  failures = 0
{.repeated section cases}
  d = LoadFromFile("py/{name}.tnet")
  expected = {py}
  if d != expected:
    print "Expected %r, got %r" % (expected, d)
    failures += 1
{.end}
  print "Python failures:", failures


def TimeDump(obj, name, n):
  start = time.time()
  for i in xrange(n):
    tnet.dumps(obj)
  elapsed = time.time() - start
  print '%s: %s dumps in %f ms' % (name, n, elapsed*1000)


def TimeDumpAll(n):
{.repeated section cases}
  obj = {py}
  TimeDump(obj, "{name}", n)
{.end}


def TimeLoad(filename, name, n):
  d = open(filename).read()
  start = time.time()
  for i in xrange(n):
    tnet.loads_prefix(d)
  elapsed = time.time() - start
  print '%s: %s loads in %f ms' % (name, n, elapsed*1000)
  return tnet.loads_prefix(d)


def TimeLoadAll(n):
{.repeated section cases}
  TimeLoad("py/{name}.tnet", "{name}", n)
{.end}


def TimeLoadAndDump(filename, n):
  obj, _ = TimeLoad(filename, "BIG", n)
  TimeDump(obj, "BIG", n)


action = sys.argv[1]
if action == 'dump':
  DumpAll()
elif action == 'load':
  LoadAll()
elif action == 'dump-perf':
  TimeDumpAll(100)
elif action == 'load-perf':
  TimeLoadAll(100)
elif action == 'perf':
  TimeLoadAndDump(sys.argv[2], int(sys.argv[3]))
else:
  raise AssertionError(action)
""")


R_TEST = jsontemplate.Template("""\
#!/usr/bin/R --vanilla --slave -f
#!/usr/bin/R -f

source("{tnet_dir}/R/tnet.R")

DumpToFile <- function(filename, data) {
  f <- file(filename, 'w')
  cat(tnet.dump(data), file=f)
  # reduce noise from diff -r
  cat('\\n', file=f)
  close(f)
}

LoadFromFile <- function(filename) {
  f <- file(filename, 'rb')  # binary
  d <- tnet.loadf(f)
  close(f)
  return(d)
}

# Dump everything
DumpAll <- function() {
{.repeated section cases}
  DumpToFile("R/{name}.tnet", {R})
{.end}
}

TimeDump <- function(obj, name, n) {
  t <- system.time(replicate(n, tnet.dump(obj)))
  cat(paste(name, ':', n, 'dumps in', t[1]*1000, 'ms of user time\\n'))
}

TimeDumpAll <- function(n) {
{.repeated section cases}
  obj <- {R}
  TimeDump(obj, "{name}", n)
{.end}
}

LoadAll <- function() {
  failures <- 0

{.repeated section cases}
  d <- LoadFromFile("R/{name}.tnet")

  cat(paste("\\t", "{name}\\n"))

  {.section r_check}
  if (!({@})) {
  {.or}
  expected <- {R}
  if (d != expected) {
  {.end}
    cat(paste("FAILED", "{name}\\n"))
    {.r_check?}
    cat("Check that failed: {r_check}\\n")
    {.or}
    cat("Expected:")
    print(expected)
    {.end}

    cat("Actual:\\n")
    print(d)
    cat('\\n\\n')

    #cat(paste("Expected", expected, "got", d, "\\n"))
    failures <- failures + 1
  }
{.end}
  cat(paste("R failures:", failures, "\\n"))
}

TimeLoad <- function(filename, test.name, n) {
  f <- file(filename)
  s <- readLines(f)[[1]]
  close(f)
  t <- system.time(
    for (i in 1:n) {
      obj <- tnet.load(s)
    })
  cat(paste(test.name, ':', n, 'loads in', t[1]*1000, 'ms of user time\\n'))
  obj  # return it so TimeDump can use it
}

TimeLoadAll <- function(n) {
{.repeated section cases}
  TimeLoad("R/{name}.tnet", "{name}", n)
{.end}
}

TimeLoadAndDump <- function(filename, n) {
  r <- TimeLoad(filename, "TimeLoadAndDump", n)
  TimeDump(r$value, "TimeLoadAndDump", n)
}

argv <- commandArgs(trailingOnly=T)
action <- argv[[1]]
if (action == 'dump') {
  DumpAll()
} else if (action == 'load') {
  LoadAll()
} else if (action == 'dump-perf') {
  TimeDumpAll(100)
} else if (action == 'load-perf') {
  TimeLoadAll(100)
} else if (action == 'perf') {
  TimeLoadAndDump(args[[2]], as.integer(args[[3]]))
} else {
  stop(action)
}
""")

WORK_DIR = '_tmp/multi_test'


def _WriteCode(filename, code_str):
  path = os.path.join(WORK_DIR, filename)
  open(path, 'w').write(code_str)
  os.chmod(path, 0700)
  print >>sys.stderr, 'Wrote %s' % path


def GenCode():
  # generate a small script for each test

  tnet_dir = os.path.dirname(os.path.abspath(sys.argv[0]))
  data = {
      'cases': _test_cases,
      'tnet_dir': tnet_dir,
      }

  py = PY_TEST.expand(data)
  _WriteCode('tnet_test_cases.py', py)

  r = R_TEST.expand(data)
  _WriteCode('tnet_test_cases.R', r)


if __name__ == '__main__':
  GenCode()
